import torch
import torchvision
import matplotlib.pyplot as plt
from classifier import Classifier

classifier = Classifier()
transform = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,))
        ])
train_dataset = torchvision.datasets.MNIST(
        './data/', train=True, download=True,
        transform=transform)
test_dataset = torchvision.datasets.MNIST(
        './data/', train=False, download=True,
        transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)

def show_examples():
    examples = enumerate(test_loader)
    batch_idx, (example_data, example_targets) = next(examples)
    print(example_targets)
    print(example_data.shape)
    fig = plt.figure()
    for i in range(6):
        plt.subplot(2,3,i+1)
        plt.tight_layout()
        plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
        plt.title("Ground Truth: {}".format(example_targets[i]))
        plt.xticks([])
        plt.yticks([])
    plt.show()


def main():
    classifier.train(train_loader)
    classifier.test(test_loader)


if __name__ == '__main__':
    main()
    # show_examples()
